use crate::curve::*;
use crate::U256;
use std::fmt;

/*
 * ECC private keys are generated using the method of extra random bits,
 * equivalent to that described in FIPS 186-4, Appendix B.4.1.
 *
 * d = (c mod(n–1))	where c is a string of random bits, 64 bits longer
 *						than requested
 * 1 <= c mod(n-1) <= n-2
 *
 * This method generates a private key uniformly distributed in the range
 * [1, n-2].
 */
// definitions for Private/Public key
#[derive(Eq, Clone, Default)]
#[cfg(feature = "mont_d")]
pub struct PrivateKey {
    priv_key: U256,
    pubx: U256,
    puby: U256,
    mont_d: U256,
    d_inv: U256,
}
#[derive(Eq, Clone, Default)]
#[cfg(not(feature = "mont_d"))]
pub struct PrivateKey {
    priv_key: U256,
    pubx: U256,
    puby: U256,
    d_inv: U256,
}

impl PartialEq for PrivateKey {
    // noneed to compare puby
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.priv_key == other.priv_key && self.pubx == other.pubx
    }
}

impl fmt::Display for PrivateKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "\nprivate key: {}\npubx: {}\npuby: {}\n",
            self.priv_key, self.pubx, self.puby
        )
    }
}

impl PrivateKey {
    pub fn new<const K0: u64, const NK0: u64, const A_IS_N3: bool>(
        curve: &dyn CurveSW<K0, NK0, A_IS_N3>,
    ) -> PrivateKey {
        let priv_key = U256::random();
        let mut res = PrivateKey {
            priv_key,
            ..Default::default()
        };
        while !res.init(curve) {
            res.priv_key = U256::random();
        }
        res.calc_pa(curve);
        res
    }
    pub fn from<const K0: u64, const NK0: u64, const A_IS_N3: bool>(
        curve: &dyn CurveSW<K0, NK0, A_IS_N3>,
        privk: &U256,
    ) -> Option<PrivateKey> {
        let priv_key = privk.clone();
        let mut res = PrivateKey {
            priv_key,
            ..Default::default()
        };
        if res.init(curve) {
            res.calc_pa(curve);
            Some(res)
        } else {
            None
        }
    }
    fn init<const K0: u64, const NK0: u64, const A_IS_N3: bool>(
        &mut self,
        curve: &dyn CurveSW<K0, NK0, A_IS_N3>,
    ) -> bool {
        let mut n1 = curve.n().clone();
        n1.usub(1);
        // mod N-1
        let d = self.priv_key.modp(&n1, 0);
        if d.is_zero() {
            return false;
        }
        // _d in [1 .. N-2]
        let mut d_inv = d.clone();
        self.priv_key = d;
        d_inv.uadd(1);
        // if d_inv == *curve.n() { return false; }
        // d_inv = (1 + d)^-1
        d_inv = d_inv.mod_inverse(curve.n());
        self.d_inv = curve.to_montgomeryN(&d_inv);
        #[cfg(feature = "mont_d")]
        {
            self.mont_d = curve.to_montgomeryN(&self.priv_key);
        }
        true
    }
    #[cfg(feature = "mont_d")]
    pub fn mont_d(&self) -> &U256 {
        &self.mont_d
    }
    // Di() return (1 + dA)^-1 in montgomery form modN
    pub fn d_inv(&self) -> &U256 {
        &self.d_inv
    }
    pub fn public_key(&self) -> (U256, U256) {
        (self.pubx.clone(), self.puby.clone())
    }
    pub fn private_key(&self) -> &U256 {
        &self.priv_key
    }
    #[inline]
    fn calc_pa<const K0: u64, const NK0: u64, const A_IS_N3: bool>(
        &mut self,
        curve: &dyn CurveSW<K0, NK0, A_IS_N3>,
    ) {
        (self.pubx, self.puby) = curve.scalar_base_mult(&self.priv_key);
    }
}

/*
 * SM2 sign/verify
 * 			e = HASH (msg)
 * 			Pa ...	Public key, Px/Py
 * 			dA ...	private key
 * 	sign
 * 			k = random   256bits   Label getk
 * 			x1, y1 = k * G
 * 			r = e + x1
 *			if r == 0 modN  goto getk
 *			s = (1 + dA)^-1 * (k - r * dA)
 *			s = (1 + dA)^-1 * (k + r) - r
 *			if s == 0 modN goto getk
 *			if r + s == 0 modN goto getk
 *			return r, s, y1_is_odd
 *
 * 	verify
 * 			t = r + s modN
 * 			if t == 0 fail
 * 			x2, y2 = s*G + t*Pa
 * 			if x2 + e == r modN success else fail
 *
 *	recover public key
 *			Psx = r - e
 *			Psy = get_Py ( Psx, Psy_is_odd )
 *			Ps ...	Psx/Psy
 *			u1 = (r + s)^-1
 *			u2 = s * u1
 *			Pa = (r+s)^1 * ( k*G - s*G) = (r+s)^-1 * Ps - s*(r+s)^-1 * G
 *			Pa = u1 * Ps - u2 * G
 */

impl PrivateKey {
    // sign return ecInd,  false for Py(of sign) even, true for odd
    pub fn sign<const K0: u64, const NK0: u64, const A_IS_N3: bool>(
        &self,
        curve: &dyn CurveSW<K0, NK0, A_IS_N3>,
        msg: &U256,
    ) -> (U256, U256, bool) {
        loop {
            let mut r: U256;
            let mut k: U256;
            let ret: bool;
            loop {
                let x1: U256;
                let y1: U256;
                (k, x1, y1) = PrivateKey::gen_keypair(curve);
                r = msg.mod_add(&x1, curve.n());
                if !r.is_zero() {
                    ret = y1.is_odd();
                    break;
                }
            }
            let s: U256;
            #[cfg(feature = "mont_d")]
            {
                let k = curve.to_montgomeryN(&k);
                let rp = curve.to_montgomeryN(&r);
                // tmp = r * dA
                let tmp = curve.n_mult(&rp, self.mont_d());
                // k = k - r * dA
                let k = k.mod_sub(&tmp, curve.n());
                // tmp = (1 + dA)^-1 * (k - r * dA)
                let tmp = curve.n_mult(self.d_inv(), &k);
                s = curve.from_montgomeryN(&tmp);
            }
            #[cfg(not(feature = "mont_d"))]
            {
                let tmp = k.mod_add(&r, curve.n());
                let tmp = curve.to_montgomeryN(&tmp);
                let tmp = curve.n_mult(&tmp, self.d_inv());
                let tmp = curve.from_montgomeryN(&tmp);
                s = tmp.mod_sub(&r, curve.n());
            }
            if s.is_zero() {
                continue;
            }
            // SM2 sign s may above halfN
            let tmp = r.mod_add(&s, curve.n());
            if !tmp.is_zero() {
                return (r, s, ret);
            }
        }
    }
    // generate random private/public keypairs
    pub fn gen_keypair<const K0: u64, const NK0: u64, const A_IS_N3: bool>(
        curve: &dyn CurveSW<K0, NK0, A_IS_N3>,
    ) -> (U256, U256, U256) {
        let mut priv_key;
        let mut n1 = curve.n().clone();
        n1.usub(1);
        loop {
            priv_key = U256::random();
            if priv_key > n1 {
                priv_key.sub_from(&n1);
            }
            if !priv_key.is_zero() {
                break;
            }
        }
        let (px, py) = curve.scalar_base_mult(&priv_key);
        (priv_key, px, py)
    }
}

// verify return bool, true for success(ok)
pub fn ec_verify<const K0: u64, const NK0: u64, const A_IS_N3: bool>(
    curve: &dyn CurveSW<K0, NK0, A_IS_N3>,
    r: &U256,
    s: &U256,
    pubx: &U256,
    puby: &U256,
    msg: &U256,
) -> bool {
    if r.is_zero() {
        return false;
    }
    if s.is_zero() {
        return false;
    }
    if *curve.n() < *r {
        return false;
    }
    if *curve.n() < *s {
        return false;
    }
    let t = r.mod_add(s, curve.n());
    if t.is_zero() {
        return false;
    }
    // q = s*G + t * Pub
    // x2, y2 = s*G + t * Pub
    let q = curve.cmult(pubx, puby, &t, s);
    if q.z.is_zero() {
        return false;
    }
    // t = r - msg
    let t = r.mod_sub(msg, curve.n());
    let zz = curve.msqr(q.z(), 1);
    if curve.prod_equal(q.x(), &t, &zz) {
        return true;
    }
    if t < curve.p_minus_n() {
        let (t, _) = t.add(curve.n());
        if curve.prod_equal(q.x(), &t, &zz) {
            return true;
        }
    }
    false
}

// vecover public key, from r/s/v/msg, v for pubY is odd
// return true for successfully recover public key
pub fn ec_recover<const K0: u64, const NK0: u64, const A_IS_N3: bool>(
    curve: &dyn CurveSW<K0, NK0, A_IS_N3>,
    r: &U256,
    s: &U256,
    msg: &U256,
    v: bool,
) -> (U256, U256, bool) {
    // p.x = r - msg
    let px = r.mod_sub(msg, curve.n());
    let py: U256 = Default::default();
    if let Some(py) = curve.y_recover(&px, v) {
        // if no verify r/s, only cost 10-20 us on Thinkpad T440s
        //return (px, py, true);
        let u1 = r.mod_add(s, curve.n());
        if u1.is_zero() {
            return (px, py, false);
        }
        // u1 = (r + s)^-1
        let u1 = u1.mod_inverse(curve.n());
        let u2 = curve.to_montgomeryN(s);
        let u1p = curve.to_montgomeryN(&u1);
        let u2 = curve.n_mult(&u2, &u1p);
        // u2 = s * u1
        let u2 = curve.from_montgomeryN(&u2);
        // 0 < u2 < N
        if u2.is_zero() {
            return (px, py, false);
        }
        // u2 = -u2 = - s* u1
        let (u2, _) = curve.n().sub(&u2);
        // Pa = u1 * Ps - u2 * G
        let q = curve.cmult(&px, &py, &u1, &u2);
        let (px, py) = curve.apply_z(&q);
        (px, py, true)
    } else {
        return (px, py, false);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use measure::Measure;

    #[test]
    fn test_new() {
        let curve = CurveSM2::new("sm2p");
        let pk = PrivateKey::new(&curve);
        println!("dump key: {}", pk);
    }

    #[test]
    fn test_sign() {
        let curve = CurveSM2::new("sm2p");
        let pk = PrivateKey::new(&curve);
        let (pubx, puby) = pk.public_key();
        let msg = U256::random();
        for _i in 0..10 {
            let (r, s, _ec_ind) = pk.sign(&curve, &msg);
            assert!(!r.is_zero());
            assert!(!s.is_zero());
            let t = r.mod_add(&s, curve.n());
            assert!(!t.is_zero());
            // q = s*G + t * Pub
            // x2, y2 = s*G + t * Pub
            let q = curve.cmult(&pubx, &puby, &t, &s);
            assert!(!q.z.is_zero());
            let (x, _y) = curve.apply_z(&q);
            let t = x.mod_add(&msg, curve.n());
            assert!(t == r);
        }
    }

    #[test]
    fn test_verify() {
        let curve = CurveSM2::new("sm2p");
        let pk = PrivateKey::new(&curve);
        let (pubx, puby) = pk.public_key();
        let msg = U256::random();
        for _i in 0..10 {
            let (r, s, _ec_ind) = pk.sign(&curve, &msg);
            assert!(!r.is_zero());
            assert!(!s.is_zero());
            assert!(ec_verify(&curve, &r, &s, &pubx, &puby, &msg));
        }
    }

    #[test]
    fn test_recover() {
        let curve = CurveSM2::new("sm2p");
        let pk = PrivateKey::new(&curve);
        let (pubx, puby) = pk.public_key();
        let msg = U256::random();
        for _i in 0..10 {
            let (r, s, ec_ind) = pk.sign(&curve, &msg);
            let (px, py, ok) = ec_recover(&curve, &r, &s, &msg, ec_ind);
            assert!(ok, "ec_verify failed");
            assert!(px == pubx);
            assert!(py == puby);
        }
    }

    #[test]
    #[ignore]
    #[allow(non_snake_case)]
    fn bench_ec_sign() {
        let curve = CurveSM2::new("sm2p");
        let pk = PrivateKey::new(&curve);
        let msg = U256::random();
        const N: u32 = 10_000;
        let mut measure = Measure::start("ec_sign bench");
        for _i in 0..N {
            let (_r, _s, _ec_ind) = pk.sign(&curve, &msg);
        }
        measure.stop();
        let ns_ops = measure.as_ns() / (N as u64);
        let nops = N as u64 * 1_000_000_000 / measure.as_ns();
        println!("ec_sign cost {} ns per Op, {} msg/sec", ns_ops, nops);
    }

    #[test]
    #[ignore]
    #[allow(non_snake_case)]
    fn bench_ec_recover() {
        let curve = CurveSM2::new("sm2p");
        let pk = PrivateKey::new(&curve);
        let msg = U256::random();
        const N: u32 = 10_000;
        let (r, s, ec_ind) = pk.sign(&curve, &msg);
        let mut measure = Measure::start("ec_recover bench");
        for _i in 0..N {
            let (_px, _py, _ok) = ec_recover(&curve, &r, &s, &msg, ec_ind);
        }
        measure.stop();
        let ns_ops = measure.as_ns() / (N as u64);
        let nops = N as u64 * 1_000_000_000 / measure.as_ns();
        println!("ec_recover cost {} ns per Op, {} msg/sec", ns_ops, nops);
    }
}
